#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@ File: evaluation.py
@ Author: YuJX
@ Reference: BBL-FasterRCNN-pytorch, https://github.com/bubbliiiing/faster-rcnn-pytorch
@ Function : 自定义的与数据集的获取、处理、组织相关部件

"""
import os
import torch
from collections import defaultdict

import numpy as np
from PIL import Image
import shutil

from tqdm import tqdm
import matplotlib
matplotlib.use('Agg')
from matplotlib import pyplot as plt

from utils.data import cvtColor, resize_image, preprocess_input, get_new_img_size
from utils.map import get_coco_map, get_map
from torch.utils.tensorboard import SummaryWriter

from models.detectors.FasterRCNN import FasterRCNNBboxDecoder

class EvalCallback():
    def __init__(self, detector, decoder, input_shape, class_names, num_classes, val_lines, log_dir, cuda, \
            map_out_path=".temp_map_out", max_boxes=100, confidence=0.05, nms_iou=0.5, letterbox_image=True, MINOVERLAP=0.5, eval_flag=True, period=1):
        super(EvalCallback, self).__init__()

        
        self.input_shape        = input_shape
        self.class_names        = class_names
        self.num_classes        = num_classes
        self.val_lines          = val_lines
        self.log_dir            = log_dir
        self.cuda               = cuda
        self.map_out_path       = map_out_path
        self.max_boxes          = max_boxes
        self.confidence         = confidence
        self.nms_iou            = nms_iou
        self.letterbox_image    = letterbox_image
        self.MINOVERLAP         = MINOVERLAP
        self.eval_flag          = eval_flag
        self.period             = period
        
        self.eval_map_record = defaultdict(list)
        self.bring_history = True
        self.writer     = SummaryWriter(self.log_dir)
        
        self.model      = detector
        self.bbox_util  = decoder

    #---------------------------------------------------#
    #   检测图片
    #---------------------------------------------------#
    def get_map_txt(self, image_id, image, class_names, map_out_path):
        f = open(os.path.join(map_out_path, "detection-results/"+image_id+".txt"),"w")
        #---------------------------------------------------#
        #   计算输入图片的高和宽
        #---------------------------------------------------#
        image_shape = np.array(np.shape(image)[0:2])
        # input_shape = get_new_img_size(image_shape[0], image_shape[1])
        #---------------------------------------------------------#
        #   在这里将图像转换成RGB图像，防止灰度图在预测时报错。
        #   代码仅仅支持RGB图像的预测，所有其它类型的图像都会转化成RGB
        #---------------------------------------------------------#
        image       = cvtColor(image)
        
        #---------------------------------------------------------#
        #   给原图像进行resize，resize到短边为600的大小上
        #---------------------------------------------------------#
        image_data  = resize_image(image, [self.input_shape[1], self.input_shape[0]])
        #---------------------------------------------------------#
        #   添加上batch_size维度
        #---------------------------------------------------------#
        image_data  = np.expand_dims(np.transpose(preprocess_input(np.array(image_data, dtype='float32')), (2, 0, 1)), 0)

        with torch.no_grad():
            images = torch.from_numpy(image_data)
            if self.cuda:
                images = images.cuda()

            self.model.eval()
            output = self.model(images)
            #-------------------------------------------------------------#
            #   对建议框进行解码，获得预测框
            #   对预测框进行NMS，获得最终结果
            #-------------------------------------------------------------#
            results = self.bbox_util.forward(output, image_shape, self.input_shape, nms_iou = self.nms_iou, confidence = self.confidence)
            #--------------------------------------#
            #   如果没有检测到物体，则返回原图
            #--------------------------------------#
            if len(results[0]) <= 0:
                return 

            top_label   = np.array(results[0][:, 5], dtype = 'int32')
            top_conf    = results[0][:, 4]
            top_boxes   = results[0][:, :4]

        top_100     = np.argsort(top_conf)[::-1][:self.max_boxes]
        top_boxes   = top_boxes[top_100]
        top_conf    = top_conf[top_100]
        top_label   = top_label[top_100]

        for i, c in list(enumerate(top_label)):
            predicted_class = self.class_names[int(c)]
            box             = top_boxes[i]
            score           = str(top_conf[i])

            top, left, bottom, right = box
            if predicted_class not in class_names:
                continue

            f.write("%s %s %s %s %s %s\n" % (predicted_class, score[:6], str(int(left)), str(int(top)), str(int(right)),str(int(bottom))))

        f.close()
        return 
    
    def on_epoch_end(self, epoch):
        if epoch % self.period == 0 and self.eval_flag:
            if not os.path.exists(self.map_out_path): # 训练时为临时文件夹，评估结束后释放
                os.makedirs(self.map_out_path)
            if not os.path.exists(os.path.join(self.map_out_path, "ground-truth")):
                os.makedirs(os.path.join(self.map_out_path, "ground-truth"))
            if not os.path.exists(os.path.join(self.map_out_path, "detection-results")):
                os.makedirs(os.path.join(self.map_out_path, "detection-results"))
                
            print("\n\033[1;36;40mGet map\033[0m")
            for annotation_line in tqdm(self.val_lines):
                line        = annotation_line.split()
                image_id    = os.path.basename(line[0]).split('.')[0]
                #------------------------------#
                #   读取图像并转换成RGB图像
                #------------------------------#
                image       = Image.open(line[0])
                #------------------------------#
                #   获得预测框
                #------------------------------#
                gt_boxes    = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
                #------------------------------#
                #   获得预测txt
                #------------------------------#
                self.get_map_txt(image_id, image, self.class_names, self.map_out_path)
                
                #------------------------------#
                #   获得真实框txt
                #------------------------------#
                with open(os.path.join(self.map_out_path, "ground-truth/"+image_id+".txt"), "w") as new_f:
                    for box in gt_boxes:
                        left, top, right, bottom, obj = box
                        obj_name = self.class_names[obj]
                        new_f.write("%s %s %s %s %s\n" % (obj_name, left, top, right, bottom))
                        
            print("Calculate Map.")
            # try:
            #     temp_map = get_coco_map(class_names = self.class_names, path = self.map_out_path)[1]
            # except:
            temp_map = get_map(self.MINOVERLAP, False, path = self.map_out_path)

            if self.bring_history:
                self.bring_history = False
                for epoch in self.eval_map_record['epoch']:
                    # 写入历史值
                    self.writer.add_scalar('eval_map', self.eval_map_record['eval_map'][epoch-1], epoch)
                    with open(os.path.join(self.log_dir, "eval_map.txt"), 'a') as f:
                        f.write(str(self.eval_map_record['eval_map'][epoch-1]))
                        f.write("\n")
                        
            self.eval_map_record['eval_map'].append(temp_map)
            self.eval_map_record['epoch'].append(epoch)

            # 写入tb
            self.writer.add_scalar('eval_map', temp_map, epoch)
            self.writer.add_text('eval_map', str(temp_map), epoch)
            # 写入txt
            with open(os.path.join(self.log_dir, "eval_map.txt"), 'a') as f:
                f.write(str(temp_map))
                f.write("\n")
            
            # 绘图
            plt.figure()
            plt.plot(self.eval_map_record['epoch'], self.eval_map_record['eval_map'], 'red', linewidth = 2, label='eval map')

            plt.grid(True)
            plt.xlabel('Epoch')
            plt.ylabel('Map %s'%str(self.MINOVERLAP))
            plt.title('A Map Curve')
            plt.legend(loc="upper right")

            plt.savefig(os.path.join(self.log_dir, "eval_map.png"))
            plt.cla()
            plt.close("all")

            print("Get map done.")
            shutil.rmtree(self.map_out_path)
